package NLP;

import com.hankcs.hanlp.mining.word2vec.Word2VecTrainer;
import com.hankcs.hanlp.mining.word2vec.WordVectorModel;
import org.junit.Test;

import java.io.IOException;
import java.util.List;
import java.util.Map;

public class Word2Vec {
    private String trainFileName = "corpus//simplified//jueju//words//segmentedResultPerPoem.txt";
    private String modelFileName = "E:\\LiuChen\\IDEA Java\\GAPoemGenerator\\corpus\\simplified\\jueju\\words\\word2vecModel.txt";
    private WordVectorModel wordVectorModel;
    /**
     * 训练模型
     */
    public void train(){
        Word2VecTrainer trainerBuilder = new Word2VecTrainer();
        WordVectorModel wordVectorModel = trainerBuilder.train(trainFileName,modelFileName);
    }

    /**
     * 加载模型
     * @return 成功加载返回true, 否则返回false
     */
    public boolean loadModel(){
        try {
            wordVectorModel = new WordVectorModel(modelFileName);
            if (wordVectorModel!=null){
                return true;
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
        return false;
    }

    /**
     * 计算两个词的相似程度
     * @param what 第一个词
     * @param with 第二个词
     * @return
     */
    public float similarity(String what,String with){
        return wordVectorModel.similarity(what,with);
    }

    /**
     * 给出最相似的size个词语
     * @param word 词
     * @param size rank top前size个词语
     * @return
     */
    public List<Map.Entry<String,Float>> nearest(String word,int size){
        return wordVectorModel.nearest(word,size);
    }
}
